package com.web709.controller;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import io.reactivex.Flowable;
import io.swagger.annotations.Api;
import io.swagger.annotations.ApiImplicitParam;
import io.swagger.annotations.ApiImplicitParams;
import io.swagger.annotations.ApiOperation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;

import java.io.IOException;
import java.io.OutputStream;
import java.util.Arrays;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
@Api(tags = "Ai接口")
@RestController
public class DashScopeController {

    private static final Logger logger = LoggerFactory.getLogger(DashScopeController.class);
    private static final ExecutorService threadPool = Executors.newCachedThreadPool();

    @GetMapping(value = "/chat-stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
@ApiImplicitParams({
        @ApiImplicitParam(name = "content", value = "用户输入", required = true),
        //@ApiImplicitParam(name = "token", required = false, paramType = "header", dataType = "String")
})
@ApiOperation(value = "Ai对话", notes = "流式对话")
public SseEmitter stream(@RequestParam String content, @RequestHeader(required = false) String token) {
    // 验证和清理用户输入
    content = validateAndSanitizeInput(content);

    // 验证 token 的有效性
    if (isValidToken(token)) {
        SseEmitter emitter = new SseEmitter();
        emitter.onTimeout(() -> emitter.complete());
        try {
            emitter.send(SseEmitter.event().data("Invalid token"));
            emitter.complete();
        } catch (IOException e) {
            emitter.completeWithError(e);
        }
        return emitter;
    }

    SseEmitter emitter = new SseEmitter();
    emitter.onTimeout(() -> emitter.complete());

    String finalContent = content;
    threadPool.submit(() -> {
        try {
            Generation gen = new Generation();
            Message userMsg = Message.builder().role(Role.USER.getValue()).content(finalContent).build();
            GenerationParam param = buildGenerationParam(userMsg);
            Flowable<GenerationResult> result = gen.streamCall(param);

            result.subscribe(
                    data -> {
                        if (data.getOutput() != null && !data.getOutput().getChoices().isEmpty()) {
                            Message message = data.getOutput().getChoices().get(0).getMessage();
                            if (message != null) {
                                try {
                                    emitter.send(message.getContent());

                                } catch (IOException e) {
                                    emitter.completeWithError(e);
                                }
                            }
                        }
                    },
                    error -> {
                        logger.error("Error during streamCall: {}", error.getMessage());
                        emitter.completeWithError(error);
                    },
                    () -> {
                        logger.info("SSE subscription completed");
                        emitter.complete();
                    }
            );
        } catch (Exception e) {
            logger.error("Exception in thread pool task: {}", e.getMessage());
            emitter.completeWithError(e);
        }
    });

    return emitter;
}

private boolean isValidToken(String token) {
    // 实现 token 验证逻辑
    return token != null && !token.trim().isEmpty(); // 示例代码，实际应替换为真实的验证逻辑
}


    private GenerationParam buildGenerationParam(Message userMsg) {
        return GenerationParam.builder()
                // 若没有配置环境变量，请用百炼API Key将下行替换为：.apiKey("sk-xxx")
                .apiKey("sk-35d0dd27d19b40698ffa29157c97ebbc")
                .model("deepseek-r1")
                .messages(Arrays.asList(userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .incrementalOutput(true)
                .build();
    }

    private String validateAndSanitizeInput(String input) {
        // 实现输入验证和清理逻辑
        if (input == null || input.trim().isEmpty()) {
            throw new IllegalArgumentException("输入不能为空");
        }
        // 这里可以使用正则表达式或其他方式确保输入是安全的
        input = input.trim();
        if (input.length() > 1000) {
            throw new IllegalArgumentException("输入长度不能超过1000个字符");
        }
        return input;
    }
}